from pathlib import Path
from rosbags.typesys import get_types_from_msg, register_types
import logging
import os

class Register(object):
    """ register custom msg types with support for both install and source modes
    """

    def register_msg_types(self, msg_path: str, use_source_mode: bool = False):
        """
        注册自定义消息类型，支持两种模式：
        1. 默认模式：从install目录读取消息（install/*_msgs/share/*_msgs/msg/）
        2. 源文件模式：从ROS2 msg源文件目录读取（通过use_source_mode=True启用）
        
        :param msg_path: 基础路径（install目录或ROS2 msg源文件根目录）
        :param use_source_mode: 是否使用源文件模式（默认False，即使用install模式）
        """
        add_types = {}
        msg_files = []
        
        if use_source_mode:
            # 源文件模式：使用原有ROS2 msg路径配置
            msg_files = self._get_source_mode_msg_files(msg_path)
            mode_name = "ROS2 msg源文件模式"
        else:
            # 默认模式：从install目录读取
            msg_files = self._get_install_mode_msg_files(msg_path)
            mode_name = "install目录模式"
        
        # 验证是否找到消息文件
        if not msg_files:
            error_msg = f"在{mode_name}下未找到任何.msg文件，路径: {msg_path}"
            logging.error(error_msg)
            raise FileNotFoundError(error_msg)
        
        # 读取并注册消息类型
        for pathstr in msg_files:
            try:
                msgpath = Path(pathstr)
                msgdef = msgpath.read_text(encoding="utf-8")
                add_types.update(get_types_from_msg(msgdef, self._guess_msgtype(msgpath)))
            except Exception as e:
                logging.warning(f"读取消息文件失败: {pathstr}, 错误: {str(e)}")
        
        if not add_types:
            raise RuntimeError("未成功解析任何消息类型，无法完成注册")
        
        register_types(add_types)
        logging.info(f"成功在{mode_name}下注册{len(msg_files)}个消息类型")
        
        # 导入消息类型（保持原有需要的类型）
        self._import_message_types()

    def _get_source_mode_msg_files(self, msg_path: str) -> list:
        """获取源文件模式下的.msg文件列表（原有逻辑）"""
        # 原有固定路径配置
        file_structure = {
            "sensors_msgs/gnss_imu_msgs/msg": [
                "Gnss.msg",
                "Imu.msg",
                "Inspva.msg",
                "GpsTime.msg",
            ],
            "sensors_msgs/uss_msgs/msg": [
                "Vhm.msg"
            ],
            "vehicle_msgs/msg": ["VehicleMotionIpd.msg"],
            "parking_ego_motion_msgs/msg": ["DrResult.msg",
                                            "ZoneTransformStamped.msg",
                                            "ZoneTransform.msg"],
            "common_msgs/msg": [
                "ComHeader.msg",
                "TimeStamp.msg",
                "Point2F.msg",
                "Point2D.msg",
                "Point3D.msg",
                "Attitude.msg",
            ],
        }
        
        # 构建完整路径
        msg_files = [
            os.path.join(msg_path, folder, file)
            for folder, files in file_structure.items()
            for file in files
        ]
        
        # 过滤不存在的文件
        existing_files = []
        for file_path in msg_files:
            if os.path.exists(file_path):
                existing_files.append(file_path)
            else:
                logging.warning(f"源文件模式下未找到消息文件: {file_path}")
        
        return existing_files

    def _get_install_mode_msg_files(self, install_path: str) -> list:
        """获取install模式下的.msg文件列表（install/*_msgs/share/*_msgs/msg/）"""
        import glob
        
        # 查找install目录下所有符合模式的.msg文件
        # 路径格式：install/*_msgs/share/*_msgs/msg/*.msg
        install_pattern = os.path.join(
            install_path, 
            "*_msgs",          # 匹配以_msgs结尾的文件夹
            "share", 
            "*_msgs",          # 匹配同名的_msgs文件夹
            "msg", 
            "*.msg"
        )
        
        msg_files = glob.glob(install_pattern, recursive=True)
        
        if not msg_files:
            logging.warning(f"install模式下未找到任何.msg文件，搜索路径: {install_pattern}")
        
        return msg_files

    def _guess_msgtype(self, path: Path) -> str:
        """Guess message type name from path."""
        name = path.relative_to(path.parents[2]).with_suffix('')
        if 'msg' not in name.parts:
            name = name.parent / 'msg' / name.name
        return str(name)

    def _import_message_types(self):
        """导入所需的消息类型（保持原有需要的类型）"""
        from rosbags.typesys.types import (
            # gnss_imu_msgs__msg__Gnss as Gnss,
            gnss_imu_msgs__msg__Imu as Imu,
            gnss_imu_msgs__msg__Inspva as Inspva,
            gnss_imu_msgs__msg__GpsTime as GpsTime,
            vehicle_msgs__msg__VehicleMotionIpd as VehicleMotionIpd,
            # hd_map_ehr_to_em_msgs__msg__EhrToEmData as EhrToEmData,
            parking_ego_motion_msgs__msg__DrResult as DrResult,
            common_msgs__msg__ComHeader as ComHeader,
            common_msgs__msg__TimeStamp as TimeStamp,
            # map_ehp_msgs__msg__PositionMsg as PositionMsg,
            cus_app_msgs__msg__FctActuationRequest as FctActuationRequest,
            cus_app_msgs__msg__ApaTarget as ApaTarget,
        )
    